import torch
from ..utils.get_config import get_cfg
import os.path as osp

cfg = get_cfg(osp.join(osp.abspath("./"), 'configs/config.yaml'))
DEVICE = cfg['DEVICE']


def get_acc(loader, model, is_validation=False):
    # return the accuracy on a dataset

    correct = 0

    model.eval()
    for data in loader:
        data = data.to(DEVICE)
        with torch.no_grad():
            emb, pred = model(data)
            pred = pred.argmax(dim=1)
            label = data.y

        if model.task == 'node':
            # get results from the validation of test dataset
            mask = data.val_mask if is_validation else data.test_mask
            pred = pred[mask]
            label = label[mask]

        correct += pred.eq(label).sum().item()

    if model.task == 'graph':
        total = len(loader.dataset)
    else:
        total = sum(torch.sum(mask).item() for _ in loader.dataset)
    return (correct/total)
